# 1. 加载环境变量
import os
from dotenv import load_dotenv
from llama_index.core.agent.workflow import ReActAgent

load_dotenv()

# 2. 初始化 LLM
from langchain.chat_models import ChatOpenAI



llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0  # 0表示输出更确定，适合需要精确决策的场景
)

# 3. 定义多工具（天气+计算器）
from langchain_core.tools import tool

# 工具 1：天气查询
@tool(
    description="用于查询中国内地城市的实时最高气温，参数为中文城市名（如'北京'，不可用英文），返回格式：成功→'城市名 气温'（如'北京 25℃'），失败→'暂无XX城市天气数据'。当需要获取城市气温时，必须调用该工具。"
)
def WeatherTool(city: str) -> str:
    weather_data = {"北京": "25℃", "上海": "28℃", "广州": "32℃", "深圳": "30℃"}
    return weather_data.get(city, f"暂无{city}天气数据")

# 工具 2：计算器
@tool(
    description="用于计算加减乘除数学运算，参数为字符串格式的数学表达式（如'28-25'），返回格式：成功→'表达式=结果'（如'28-25=3'），失败→'计算失败：原因'。当需要数学计算时，必须调用该工具。"
)
def CalculatorTool(expression: str) -> str:
    try:
        result = eval(expression)
        return f"{expression}={result}"
    except Exception as e:
        return f"计算失败：{str(e)}"

tools = [WeatherTool, CalculatorTool]  # 多工具列表

# 4. 定义多工具 ReAct Prompt（关键：增加异常处理规则）
from langchain_core.prompts import PromptTemplate

react_prompt_template = """
你是一个擅长解决多步骤问题的助手，会通过“思考-行动-观察”循环完成任务。你拥有以下工具：
{tools}

请严格遵守以下流程、格式和规则：
### 流程：
1. 思考（Thought）：分析当前问题、已有的观察结果，说明“下一步要做什么”以及“为什么”；
2. 行动（Action）：按格式调用工具（需工具时）；
3. 观察（Observation）：接收工具结果；
4. 终止：无需工具时输出“Final Answer”。

### 格式：
- 行动格式（必须严格遵守）：
  Action: 工具名
  Action Input: 工具参数
- 终止格式：Final Answer: 你的回答

### 特殊规则：
1. 若工具返回“暂无XX城市天气数据”，无需继续调用其他工具，直接告知用户该城市无数据；
2. 若工具返回“计算失败”，无需重试，直接告知用户错误原因；
3. 计算温差时，需用“高气温-低气温”（确保结果为正数）。

用户的问题：{input}
"""

prompt = PromptTemplate(
    template=react_prompt_template,
    input_variables=["tools", "input"],
    partial_variables={
        "tools": "\n".join([f"- {tool.name}：{tool.description}" for tool in tools])
    }
)

# 5. 创建 ReActAgent 与 AgentExecutor（多工具需增加迭代次数）
from langchain.agents import initialize_agent, AgentType, AgentExecutor

agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.REACT_DOCSTORE,
    prompt=prompt,
    verbose=True
)




agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,  # 多工具任务：最多5轮循环（避免无限重试）
    early_stopping_method="force"
)

# 6. 测试多工具协作（正常场景：北京+上海；异常场景：北京+杭州）
if __name__ == "__main__":
    # 测试 1：正常场景（北京+上海，均有数据）
    print("===== 测试正常场景：北京和上海的温差 =====")
    result1 = agent_executor.invoke({"input": "查询北京和上海的实时最高温，计算两地的温差是多少？"})
    print("正常场景结果：", result1["output"], "\n")

    # 测试 2：异常场景（北京+杭州，杭州无数据）
    print("===== 测试异常场景：北京和杭州的温差 =====")
    result2 = agent_executor.invoke({"input": "查询北京和杭州的实时最高温，计算两地的温差是多少？"})
    print("异常场景结果：", result2["output"])